from typing import Optional

import aiomysql
import logging
import os
import json
import re

# 从环境变量获取配置
DB_HOST = os.getenv("DB_HOST", "123.249.68.26")
DB_PORT = int(os.getenv("DB_PORT", "3306"))
DB_USER = os.getenv("DB_USER", "root")
DB_PASSWORD = os.getenv("DB_PASSWORD", "wxd0808wxd")
DB_NAME = os.getenv("DB_NAME", "mcpserver")


class DatabaseConnector:
    """MySQL数据库连接器"""

    def __init__(self):
        self.pool = None
        self.connected = False
        self.logger = logging.getLogger("db-connector")

    async def connect(self):
        """连接数据库"""
        try:
            self.pool = await aiomysql.create_pool(
                host=DB_HOST,
                port=DB_PORT,
                user=DB_USER,
                password=DB_PASSWORD,
                db=DB_NAME,
                minsize=1,
                maxsize=5,
                autocommit=True
            )
            self.connected = True
            self.logger.info("Database connector initialized")
        except Exception as e:
            self.logger.error(f"Database connection failed: {e}")
            self.connected = False
            raise

    async def close(self):
        """关闭连接"""
        if self.pool:
            self.pool.close()
            await self.pool.wait_closed()
            self.connected = False
            self.logger.info("Database connection closed")

    def is_connected(self) -> bool:
        return self.connected

    async def get_db_info(self, db_name: str) -> dict:
        """获取数据库信息"""
        if not self.connected:
            raise RuntimeError("Database not connected")

        async with self.pool.acquire() as conn:
            async with conn.cursor(aiomysql.DictCursor) as cur:
                await cur.execute(
                    "SELECT * FROM db_info WHERE db_name = %s",
                    (db_name)
                )
                result = await cur.fetchone()
                if not result:
                    raise ValueError(f"Database not found: {db_name}")
                return result

    async def get_instance_info(self, instance_id: str) -> dict:
        """获取实例信息"""
        if not self.connected:
            raise RuntimeError("Database not connected")

        async with self.pool.acquire() as conn:
            async with conn.cursor(aiomysql.DictCursor) as cur:
                await cur.execute(
                    "SELECT * FROM db_info WHERE instance_id = %s",
                    (instance_id,)
                )
                result = await cur.fetchone()
                if not result:
                    raise ValueError(f"Instance not found: {instance_id}")
                return result

    async def extract_host_ip(self, instance_id: str) -> Optional[str]:
        """尝试从实例ID中提取主机IP"""
        # 尝试匹配IP:port格式
        ip_match = re.match(r'(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}):\d+', instance_id)
        if ip_match:
            return ip_match.group(1)

        # 如果实例ID是域名格式，尝试查询数据库获取主机IP
        if not self.connected:
            return None
            
        async with self.pool.acquire() as conn:
            async with conn.cursor(aiomysql.DictCursor) as cur:
                await cur.execute(
                    "SELECT host_ip FROM db_info WHERE instance_id = %s",
                    (instance_id,)
                )
                result = await cur.fetchone()
                return result["host_ip"] if result else None